#!/usr/bin/env python3
import rospy
import cv2
import numpy as np
from sensor_msgs.msg import Image, CameraInfo
from geometry_msgs.msg import Pose
from cv_bridge import CvBridge
from tf.transformations import quaternion_from_matrix

class RealsenseChessboardDetector:
    def __init__(self):
        rospy.init_node('realsense_chessboard_detector')
        
        # 从参数服务器获取配置
        self.chessboard_rows = rospy.get_param('~rows', 12)
        self.chessboard_cols = rospy.get_param('~cols', 9)
        self.square_size = rospy.get_param('~markerSize', 0.03)
        self.image_is_rectified = rospy.get_param('~image_is_rectified', True)
        
        # 初始化相机参数
        self.camera_matrix = None
        self.dist_coeffs = None
        self.has_camera_info = False
        
        # 初始化ROS通信
        self.bridge = CvBridge()
        self.setup_topics()
        self.pose_pub = rospy.Publisher('~pose', Pose, queue_size=10)
        
        # 生成棋盘格3D坐标
        self.obj_points = self.generate_object_points()
        
    def setup_topics(self):
        # 订阅图像和相机信息（带remap支持）
        self.camera_info_sub = rospy.Subscriber(
            'camera_info', CameraInfo, self.camera_info_callback, queue_size=1)
        
        self.image_sub = rospy.Subscriber(
            'image', Image, self.image_callback, queue_size=1, buff_size=2**24)
        
    def generate_object_points(self):
        # 生成棋盘格3D坐标 (Z=0)
        objp = np.zeros((self.chessboard_rows*self.chessboard_cols,3), np.float32)
        objp[:,:2] = np.mgrid[0:self.chessboard_rows,0:self.chessboard_cols].T.reshape(-1,2)
        return objp * self.square_size
        
    def camera_info_callback(self, msg):
        # 解析相机内参
        self.camera_matrix = np.array(msg.K).reshape(3,3)
        self.dist_coeffs = np.array(msg.D) if not self.image_is_rectified else np.zeros(5)
        self.has_camera_info = True
        self.camera_info_sub.unregister()  # 只需获取一次
        
    def image_callback(self, msg):
        if not self.has_camera_info:
            rospy.logwarn_throttle(5, "等待相机参数...")
            return
            
        try:
            # 转换图像为OpenCV格式
            frame = self.bridge.imgmsg_to_cv2(msg, desired_encoding="bgr8")
            gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
            
            # 检测棋盘格角点
            ret, corners = cv2.findChessboardCorners(
                gray, (self.chessboard_rows, self.chessboard_cols),
                flags=cv2.CALIB_CB_ADAPTIVE_THRESH + cv2.CALIB_CB_NORMALIZE_IMAGE
            )
            
            if ret:
                # 亚像素优化
                criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 0.001)
                corners_refined = cv2.cornerSubPix(gray, corners, (11,11), (-1,-1), criteria)
                
                # 计算位姿
                ret, rvec, tvec = cv2.solvePnP(
                    self.obj_points, corners_refined,
                    self.camera_matrix, self.dist_coeffs
                )
                
                if ret:
                    # 转换旋转向量到四元数
                    R, _ = cv2.Rodrigues(rvec)
                    T = np.eye(4)
                    T[:3, :3] = R
                    T[:3, 3] = tvec.flatten()
                    quat = quaternion_from_matrix(T)
                    
                    # 发布位姿
                    pose_msg = Pose()
                    pose_msg.position.x = tvec[0][0]
                    pose_msg.position.y = tvec[1][0]
                    pose_msg.position.z = tvec[2][0]
                    pose_msg.orientation.x = quat[0]
                    pose_msg.orientation.y = quat[1]
                    pose_msg.orientation.z = quat[2]
                    pose_msg.orientation.w = quat[3]
                    self.pose_pub.publish(pose_msg)
                    
                    # 可视化（可选）
                    cv2.drawChessboardCorners(frame, (self.chessboard_rows, self.chessboard_cols), corners_refined, ret)
                    cv2.imshow("Chessboard Detection", frame)
                    cv2.waitKey(1)
                    
        except Exception as e:
            rospy.logerr("图像处理错误: %s", str(e))

if __name__ == '__main__':
    detector = RealsenseChessboardDetector()
    rospy.spin()